import os
import sys
import json
import time
import signal
import webbrowser
import multiprocessing

import tornado.ioloop
import tornado.web
import tornado.websocket

# fix Ctrl+C crash of scipy.stats on Windows
os.environ['FOR_DISABLE_CONSOLE_CTRL_HANDLER'] = '1'

from raw.train import main as raw_main
from attack.select_candidates import main as select_main
from attack.generate import main as attack_main
from defense.defend import main as defense_main
from evaluation.attack_evaluation import main as attack_evaluate
from evaluation.defense_evaluation import main as defense_evaluate
from evaluation.security_evaluation import main as security_evaluate
from attack_config import attack_config
from defense_config import defense_config

class QuestProcess(multiprocessing.Process):
    def __init__(self, quest, extra, pipe_send):
        super().__init__()
        self.quest = quest
        self.extra = extra
        self.pipe_send = pipe_send

    def run(self):
        # 忽略Ctrl+C信号，由父进程终止任务进程
        def sigint_handler(signum, frame):
            pass
        signal.signal(signal.SIGINT, sigint_handler)
        # 将stdout连接到管道
        class PipeWriter:
            def __init__(self, pipe):
                self.pipe = pipe
            def write(self, s):
                if not s.startswith('\r') and s != ' ':
                    self.pipe.send(s)
            def flush(self):
                pass
        sys.stdout = PipeWriter(self.pipe_send)
        # 执行任务
        try:
            print('started')
            print('quest:')
            print(json.dumps(self.quest, indent = 2))
            print('extra:')
            print(json.dumps(self.extra, indent = 2))
            print('log:')
            name = self.quest['name']
            qtype = self.quest['type']
            result = None
            if qtype == 0:
                raw_main(name)
            elif qtype == 1:
                select_main(name)
            elif qtype == 2:
                attack_main(name, self.quest['method'], self.quest['config'],
                            suffix = f'#{self.quest["id"]}')
            elif qtype == 3:
                defense_main(name, self.quest['method'], self.quest['config'],
                            suffix = f'#{self.quest["id"]}')
            elif qtype == 4:
                result = []
                if self.quest['evalType'] == 0:
                    for attack in self.extra['attacks']:
                        result.append(attack_evaluate(name, attack['method'], attack['config'],
                                                        suffix = f'#{attack["id"]}'))
                elif self.quest['evalType'] == 1:
                    for defense in self.extra['defenses']:
                        result.append(defense_evaluate(name, defense['method'], defense['config'],
                                                        suffix = f'#{defense["id"]}'))
                elif self.quest['evalType'] == 2:
                    defense_names = [defense['method'] for defense in self.extra['defenses']]
                    defense_configs = [defense['config'] for defense in self.extra['defenses']]
                    defense_suffixes = [f'#{defense["id"]}' for defense in self.extra['defenses']]
                    for attack in self.extra['attacks']:
                        result.append(security_evaluate(name, attack['method'], defense_names,
                                                        attack['config'], defense_configs,
                                                        attack_suffix = f'#{attack["id"]}',
                                                        defense_suffix_list = defense_suffixes))
            print('done')
        except Exception as e:
            self.pipe_send.send({ 'type': 'error', 'message': str(e) })
        else:
            self.pipe_send.send({ 'type': 'result', 'result': result })

if __name__ == '__main__':
    clients = set()
    quest_queue = []
    quests = {}
    next_id = 1
    quest_process = None
    log_pipe = None
    stopping = False

    def broadcast(msg):
        for client in clients:
            try:
                client.write_message(msg)
            except Exception as e:
                print(e)

    def start_process(quest):
        global quest_process, log_pipe
        extra = {}
        if quest['type'] == 4:
            if quest['evalType'] == 0 or quest['evalType'] == 2:
                extra['attacks'] = []
                ids = quest['evalAttacks']
                for attack_id in ids:
                    assert quests[attack_id]['type'] == 2, f'Quest #{attack_id} must be Type Attack (2)'
                    assert quests[attack_id]['status'] == 3, f'Quest #{attack_id} is not ready'
                    copy = quests[attack_id].copy()
                    copy['config'] = copy['config'].copy()
                    extra['attacks'].append(copy)
                for obj in quest['evalAttackChanges']:
                    if obj['id'] in ids:
                        extra['attacks'][ids.index(obj['id'])]['config'][obj['key']] = obj['value']
            if quest['evalType'] == 1 or quest['evalType'] == 2:
                extra['defenses'] = []
                ids = quest['evalDefenses']
                for defense_id in ids:
                    assert quests[defense_id]['type'] == 3, f'Quest #{defense_id} must be Type Defense (3)'
                    assert quests[defense_id]['status'] == 3, f'Quest #{defense_id} is not ready'
                    copy = quests[defense_id].copy()
                    copy['config'] = copy['config'].copy()
                    extra['defenses'].append(copy)
                for obj in quest['evalDefenseChanges']:
                    if obj['id'] in ids:
                        extra['defenses'][ids.index(obj['id'])]['config'][obj['key']] = obj['value']
        log_pipe, log_pipe_send = multiprocessing.Pipe(duplex = False)
        quest_process = QuestProcess(quest, extra, log_pipe_send)
        quest_process.start()
        loop.add_callback(poll_log)

    def stop_process():
        global quest_process
        quest_process.terminate()
        end_quest(error = '进程被用户终止')

    def end_quest(result = None, error = None, join = True):
        global quest_process, log_pipe
        if join:
            quest_process.join()
        quest = quest_queue.pop(0)
        if error == None:
            quest['status'] = 3
            quest['result'] = result
            broadcast({ 'type': 'stop', 'id': quest['id'], 'result': result })
        else:
            quest['status'] = 4
            quest['error'] = error
            broadcast({ 'type': 'error', 'id': quest['id'], 'error': error })
        if len(quest_queue) > 0: # 队列中还有等待中任务
            quest = quest_queue[0]
            try:
                start_process(quest) # 开启进程执行任务
            except Exception as e:
                end_quest(error = str(e), join = False)
                return
            quest['status'] = 2 # 状态设置为进行中
            broadcast({ 'type': 'start', 'id': quest['id'] })
        else:
            quest_process = None
            log_pipe = None

    def poll_log():
        global log_pipe, stopping
        if not log_pipe or stopping:
            return
        try:
            while log_pipe.poll():
                obj = log_pipe.recv()
                if type(obj) == str:
                    broadcast({ 'type': 'log', 'id': quest_queue[0]['id'], 'message': obj })
                elif type(obj) == dict:
                    if obj['type'] == 'error':
                        end_quest(error = obj['message'])
                        return
                    elif obj['type'] == 'result':
                        end_quest(result = obj['result'])
                        return
                else:
                    print(f'Unsupported pipe message type: {type(obj)}')
        except Exception as e:
            print(e)
            return
        loop.call_later(0.5, poll_log)

    def stop_server():
        global quest_process, stopping
        if stopping:
            return
        if quest_process:
            print('Stopping quest process...')
            quest_process.terminate()
            quest_process.join()
        print('Stopping server...')
        loop.stop()

    class WSHandler(tornado.websocket.WebSocketHandler):
        def open(self):
            clients.add(self)
            models = []
            for name in os.listdir('raw/models'):
                if name.endswith('.py'):
                    models.append(name[:-3])
            self.write_message({
                'type': 'config',
                'models': models,
                'attacks': attack_config,
                'defenses': defense_config
            })
            all_quests = quest_queue.copy()
            for id in quests:
                quest = quests[id]
                if quest['status'] != 1 and quest['status'] != 2:
                    all_quests.append(quest)
            self.write_message({ 'type': 'list', 'quests': all_quests })

        def on_message(self, msg):
            global next_id, stopping
            if stopping:
                return
            try:
                obj = json.loads(msg)
                command = obj['command']
                if command == 'create':
                    quest = obj['quest']
                    quest['time'] = time.strftime('%Y-%m-%d %H:%M')
                    quest['id'] = next_id
                    quest['status'] = 0
                    quests[next_id] = quest
                    next_id += 1
                    broadcast({ 'type': 'create', 'quest': quest })
                elif command == 'start':
                    quest = quests[obj['id']]
                    # 只有未开始和已中断的任务可以开始
                    assert quest['status'] in (0, 4), f'start: status must be 0 or 4, got {quest["status"]}'
                    quest_queue.append(quest)
                    if len(quest_queue) == 1: # 是队列中的第一个任务
                        try:
                            start_process(quest) # 开启进程执行任务
                        except Exception as e:
                            end_quest(error = str(e), join = False)
                            return
                        quest['status'] = 2 # 状态设置为进行中
                        broadcast({ 'type': 'start', 'id': quest['id'] })
                    else: # 排队
                        quest['status'] = 1 # 状态设置为等待中
                        broadcast({ 'type': 'queue', 'id': quest['id'] })
                elif command == 'stop':
                    quest = quests[obj['id']]
                    # 只有已开始的任务可以终止
                    assert quest['status'] == 2, f'stop: status must be 2, got {quest["status"]}'
                    stop_process()
                elif command == 'cancel':
                    quest = quests[obj['id']]
                    # 只有等待中的任务可以取消等待
                    assert quest['status'] == 1, f'cancel: status must be 1, got {quest["status"]}'
                    quest_queue.remove(quest)
                    quest['status'] = 0
                    broadcast({ 'type': 'cancel', 'id': quest['id'] })
                elif command == 'delete':
                    quest = quests[obj['id']]
                    # 只有未开始、已中断和已完成的任务可以删除
                    assert quest['status'] in (0, 3, 4), f'delete: status must be 0, 3, or 4, got {quest["status"]}'
                    del quests[obj['id']]
                    broadcast({ 'type': 'delete', 'id': quest['id'] })
                else:
                    print(f'{command}: unknown command')
            except Exception as e:
                print(e)

        def on_close(self):
            clients.remove(self)

        def check_origin(self, origin):
            return True

    # 读取文件
    data_path = 'ui_data/data.json'
    if os.path.isfile(data_path):
        try:
            with open(data_path, 'r') as f:
                obj = json.loads(f.read())
                for quest in obj['quests']:
                    quests[quest['id']] = quest
                next_id = obj['nextId']
        except Exception as e:
            print(e)
            print('Failed to load config, exiting')
            exit(1)

    # 创建web服务器
    app = tornado.web.Application([
        ('/ws', WSHandler),
        ('/(.*)', tornado.web.StaticFileHandler, { 'path': 'ui_web', 'default_filename': 'index.html' })
    ])
    hostname = 'localhost'
    port = 8888
    app.listen(port, hostname)

    loop = tornado.ioloop.IOLoop.current()

    # 处理Ctrl+C
    def sigint_handler(signum, frame):
        loop.add_callback_from_signal(stop_server)
    signal.signal(signal.SIGINT, sigint_handler)

    print(f'UI server started on http://{hostname}:{port}/')
    print('Press Ctrl+C to save and quit')

    # 打开浏览器
    try:
        webbrowser.open(f'http://{hostname}:{port}/')
    except Exception:
        pass

    # 开始循环
    try:
        loop.start()
    except Exception as e:
        print(e)

    # 保存到文件
    print('Saving config...')
    quest_list = []
    for id in quests:
        quest = quests[id]
        if quest['status'] in (1, 2):
            quest['status'] = 0
        quest_list.append(quest)
    try:
        obj = {
            "quests": quest_list,
            "nextId": next_id
        }
        json_obj = json.dumps(obj, separators = (',', ':'))
        with open(data_path, 'w') as f:
            print(json_obj, file = f)
    except Exception as e:
        print(e)

    print('bye')
